import numpy as np
from scipy.spatial.transform import Rotation as R

class RotationLimiter:
    def __init__(self, initial_rotation, max_angle_x_deg=1, max_angle_y_deg = 1, max_angle_z_deg=10):
        """
        初始化限制器，记录初始姿态，并设置 X、Y 和 Z 方向的最大旋转角度限制。

        :param initial_rotation: 初始旋转，表示为旋转矢量
        :param max_angle_x_deg: X 方向的最大旋转角度（单位：度）
        :param max_angle_y_deg: Y 方向的最大旋转角度（单位：度）
        :param max_angle_z_deg: Z 方向的最大旋转角度（单位：度）
        """
        self.initial_rotation = R.from_rotvec(initial_rotation)
        self.max_angle_x = np.radians(max_angle_x_deg)
        self.max_angle_y = np.radians(max_angle_y_deg)
        self.max_angle_z = np.radians(max_angle_z_deg)

    def update_limit(self,max_angle_x_deg=None,max_angle_y_deg=None,max_angle_z_deg=None):
        '''
        更新限制范围

        Args:
            max_angle_x_deg: X 方向的最大旋转角度（单位：度）
            max_angle_y_deg: Y 方向的最大旋转角度（单位：度）
            max_angle_z_deg: Z 方向的最大旋转角度（单位：度）

        Returns:None

        '''
        if max_angle_x_deg is not None:
            self.max_angle_x = np.radians(max_angle_x_deg)
        if max_angle_y_deg is not None:
            self.max_angle_y = np.radians(max_angle_y_deg)
        if max_angle_z_deg is not None:
            self.max_angle_z = np.radians(max_angle_z_deg)

    def apply_limits(self, current_rotation_vec):
        """
        对当前旋转进行限制，使其相对初始姿态的旋转角度不超过最大限制。

        :param current_rotation_vec: 当前旋转，表示为旋转矢量
        :return: 受限后的旋转矢量
        """
        current_rotation = R.from_rotvec(current_rotation_vec)
        
        # 计算相对旋转
        relative_rotation = self.initial_rotation.inv() * current_rotation
        relative_rotvec = relative_rotation.as_rotvec()
        
        # 分解相对旋转矢量
        x, y, z = relative_rotvec

        # 如果超出限制则进行缩放
        if abs(x) > self.max_angle_x:
            scale_x = abs(self.max_angle_x / x)
            x *= scale_x

        if abs(y) > self.max_angle_y:
            scale_y = abs(self.max_angle_y / y)
            y *= scale_y

        if abs(z) > self.max_angle_z:
            scale_z = abs(self.max_angle_z / z)
            z *= scale_z

        # 计算受限后的旋转矢量
        limited_relative_rotvec = np.array([x, y, z])
        
        # 将受限的相对旋转应用回初始姿态
        limited_rotation = self.initial_rotation*R.from_rotvec(limited_relative_rotvec)
        return limited_rotation.as_rotvec()


if __name__ == "__main__":
    # 示例使用
    initial_rotation = np.array([0.0, 0.0, 1.57])  # 初始姿态为零旋转
    limiter = RotationLimiter(initial_rotation)

    # # 当前旋转矢量，假设机器人产生了一定旋转
    # current_rotation_vec = np.array([0.02, 0.01, 0.3])  # 以旋转矢量表示
    #
    # # 应用限制
    # limited_rotation_vec = limiter.apply_limits(current_rotation_vec)
    # print("受限后的旋转矢量:", limited_rotation_vec)

    delata = np.array([0.01,0.01,-0.01])
    current_vec = initial_rotation.copy()

    for _ in range(100):
        current_pose = R.from_rotvec(current_vec)
        delta_pose = R.from_rotvec(delata)
        new_pose = (current_pose*delta_pose).as_rotvec()
        current_vec = limiter.apply_limits(new_pose)
        print(current_vec)

